# -*- coding: utf-8 -*-
import rpy2.robjects as ro
from rpy2.robjects.packages import importr #for rpy2 version 2.1 or above
import numpy as np
import sys
import rpy2.robjects.numpy2ri #automatically converts numpy arrays to R types if everything works.
from scipy.stats import poisson,norm

getmethod = ro.baseenv.get("getMethod")
StrVector = ro.StrVector

class pcaRes(ro.methods.RS4):
    """ Reflection of the S4 class 'pcares'. """
    def _loadings_get(self):
        return self.do_slot("loadings")
    def _scores_get(self):
        return self.do_slot("scores")

    loadings=property(_loadings_get,None,None, "Get loadings")
    scores=property(_scores_get,None,None, "Get loadings")

def pca(data,weight=True):
    if weight:
        # Weighting described in detail in Keenan and Kotula,
        # Surface Interface Analysis, 2004.
        aG=np.sqrt(np.sum(data,axis=1)/data.shape[1])
        bH=np.sqrt(np.sum(data,axis=0)/data.shape[0])
        aG[aG==0]=0.000000000001  # bad hack to avoid division by 0.
        bH[bH==0]=0.000000000001  # Sets all zero elements to 1E-12
        # This newaxis stuff is some numpy voodoo to avoid creating the huge
        # diagonal matrices for this matrix multiplication.  More info here:
        # http://www.mail-archive.com/numpy-discussion@scipy.org/msg02063.html
        data=(((1/aG)[:,np.newaxis])*data)*(1/bH)
        #data=bH[:,np.newaxis]*(np.dot(data,np.diag(aG,0)))
    res=ro.r.prcomp(data,retx=True,center=False)
    factors=np.array(res.rx2("rotation"))
    scores=np.array(res.rx2("x"))
    eigenvalues=np.array(res.rx2("sdev"))
    if weight:
        factors=bH[:,np.newaxis]*factors
        scores=aG[:,np.newaxis]*scores
    return factors,scores,eigenvalues

def gauss(val,mu,sigmasq):
    res=1/(np.sqrt(2*np.pi*sigmasq))*np.exp(-(val-mu)**2/(2*sigmasq))
    return res

def likelihood(vals,q,func='gauss'):
    #vals=(vals-np.min(vals))/np.max(vals-np.min(vals))
    muleft=np.average(vals[:q])
    muright=np.average(vals[q:])
    if func is 'gauss':
        varl=np.var(vals[:q])
        varr=np.var(vals[q:])
        p=len(vals)
        sigmasq=((q-1)*varl**2+(p-q-1)*varr**2)/(p-2)
        lhs=np.sum(np.array([gauss(val,muleft,sigmasq) for val in vals[:q]]))
        rhs=np.sum(np.array([gauss(val,muright,sigmasq) for val in vals[q:]]))
        #lhs=np.sum(np.array([norm.pdf(val,muleft,sigmasq) for val in vals[:q]]))
        #rhs=np.sum(np.array([norm.pdf(val,muleft,sigmasq) for val in vals[q:]]))
    elif func is 'poiss':
        lhs=np.sum(np.array([poisson.pmf(val,muright) for val in vals[q:]]))
        rhs=np.sum(np.array([poisson.pmf(val,muright) for val in vals[:q]]))
    return np.log10(lhs)+np.log10(rhs)

def screeML(vals,retls=False):
    ls=np.array([likelihood(vals,q) for q in xrange(2,len(vals)-1)])
    if retls:
        return ls
    else:
        return np.argmax(ls)+1

def reconstruct(factors,scores,nFactors=None):
    if type(nFactors).__name__=='list':
      tfactors=np.zeros((factors.shape[0],len(nFactors)))
      tscores=np.zeros((len(nFactors),scores.shape[1]))
      for i in xrange(len(nFactors)):
        tfactors[:,i]=factors[:,nFactors[i]]
        tscores[i,:]=scores[nFactors[i],:]
      return np.dot(tfactors,tscores)
    elif not nFactors:
      return np.dot(factors,scores)
    else:
      return np.dot(factors[:,:nFactors],scores[:nFactors,:])

def linfit(vals,end,windowsize=10):
    slopes=[np.polyfit(np.arange(windowsize),vals[i-windowsize:i],1)[0] for i in np.arange(end,end-10,-1)]
    avg=[np.average(slopes)]
    stdev=[np.std(slopes)]
    i=end-10
    while i>windowsize:
        slopes.append(np.polyfit(np.arange(windowsize),vals[i-windowsize:i],1)[0])
        avg.append(np.average(slopes))
        stdev.append(np.std(slopes))
        i=i-1
    return np.argmin(np.array(stdev))+windowsize

def lstsq_project(factors,data):
    prj_data=np.zeros((factors.shape[1],data.shape[1]))
    for col in xrange(data.shape[1]):
        prj=np.linalg.lstsq(factors,data[:,col])
        prj_data[:,col]=prj[0]
    return prj_data

def bpca(data,nFactors):
    """
    Bayesian PCA, not orthogonal.  Slow.
    """
    ro.r.library('pcaMethods')
    res=ro.r.bpca(data,nPcs=nFactors)
    factors=res.do_slot('scores')
    scores=res.do_slot('loadings')
    return np.asarray(factors),np.asarray(scores)    

def ica(data,nFactors):
    ro.r.library('fastICA')
    res = ro.r.fastICA(data,nFactors,method='C')
    factors=res.rx2("S")
    scores=res.rx2("A")
    return np.array(factors),np.array(scores)

def OrthoRotation(data,method='varimax'):
    ro.r.library('GPArotation')
    data=data.T  # You have to have more observations (columns) than variables
    # (rows).  If you use a row-centric rotation method, then this transpose
    # should be OK...
    orthoMethods=['targetT','pstT','entropy','quartimax','varimax','bentlerT','tandemI','tandemII','geominT','infomaxT','mccammon']
    if not method in orthoMethods:
        print 'Method %s is not intended for orthogonal application, or not implemented.  Valid methods are: '%method
        print orthoMethods
        sys.exit()
    res=ro.r.GPForth(data,method=method)
    factors=res.rx2("loadings")
    scores=res.rx2("Th")
    return np.asarray(factors).T,np.asarray(scores).T
    
def ObliqueRotation(data,method='oblimin'):
    ro.r.library('GPArotation')
    obliqueMethods=['oblimin','quartimin','targetQ','pstQ','oblimax','simplimax','bentlerQ','geominQ','cfQ','infomaxQ']
    if not method in obliqueMethods:
        print 'Method %s is not intended for oblique application, or not implemented.  Valid methods are: '%method
        print obliqueMethods
        sys.exit()
    
    res=ro.r.GPFoblq(data,method=method,normalize=True)
    factors=res.rx2("loadings")
    scores=res.rx2("Th")
    covariance=res.rx2("Phi")
    return np.asarray(factors).T,np.asarray(scores).T,np.asarray(covariance).T
    
if __name__=='__main__':
    ro.r.data('smoke',package='ca')
    ro.r.data('Harman',package='GPArotation')
    testdata=np.asarray(ro.r.Harman8)
    smokedata=np.asarray(ro.r['smoke'],dtype=np.float64)
    pcavec,pcaval=pca(testdata)
    print pcaval.shape
    # print ca(smokedata) # infinite or missing values in SVD
    print bpca(testdata,2)
    print ica(testdata,2)
    print OrthoRotation(testdata)
    print ObliqueRotation(testdata)
